import pyaudio
import asyncio
import threading
import numpy as np

class AudioStreamer:
    def __init__(self, chunk_duration=0.1, sample_rate=16000, channels=1, queue_maxsize=100):
        self.chunk_size = int(sample_rate * chunk_duration)
        self.sample_rate = sample_rate
        self.channels = channels
        self.audio_queue = asyncio.Queue(maxsize=queue_maxsize)
        self.is_recording = False
        self.pyaudio_instance = pyaudio.PyAudio()
        self.recording_thread = None
        self.playback_thread = None
        self.is_playing = False
        
    def _audio_callback(self, in_data, frame_count, time_info, status):
        """音频采集回调函数"""
        if self.is_recording:
            # 将采集到的音频数据放入队列，同时回传到扬声器
            audio_data = np.frombuffer(in_data, dtype=np.int16)
            
            # 使用线程安全的方式将音频数据放入队列
            try:
                # 获取当前事件循环（如果存在）
                loop = asyncio.get_running_loop()
                # 使用 call_soon_threadsafe 在线程安全的环境中放入队列
                loop.call_soon_threadsafe(self.audio_queue.put_nowait, in_data)
            except RuntimeError:
                # 如果没有运行的事件循环，直接放入队列（可能会在后续处理）
                try:
                    self.audio_queue.put_nowait(in_data)
                except asyncio.QueueFull:
                    # 如果队列满了，丢弃最旧的数据
                    try:
                        self.audio_queue.get_nowait()
                        self.audio_queue.put_nowait(in_data)
                    except (asyncio.QueueEmpty, asyncio.QueueFull):
                        pass
            
            # 返回音频数据以实现回传（监听）
            return (in_data, pyaudio.paContinue)
        return (None, pyaudio.paContinue)
    
    def _recording_worker(self):
        """音频采集工作线程"""
        try:
            self.stream = self.pyaudio_instance.open(
                format=pyaudio.paInt16,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                output=True,  # 启用输出以实现回传
                frames_per_buffer=self.chunk_size,
                stream_callback=self._audio_callback
            )
            
            self.stream.start_stream()
            
            while self.is_recording and self.stream.is_active():
                threading.Event().wait(0.1)
                
        except Exception as e:
            print(f"音频采集错误: {e}")
        finally:
            if hasattr(self, 'stream'):
                self.stream.stop_stream()
                self.stream.close()
    
    async def start_recording(self):
        """异步启动音频采集和回传"""
        if not self.is_recording:
            self.is_recording = True
            self.recording_thread = threading.Thread(target=self._recording_worker)
            self.recording_thread.daemon = True
            self.recording_thread.start()
            print("音频采集已启动，音频将通过扬声器回传...")
    
    async def stop_recording(self):
        """停止音频采集"""
        if self.is_recording:
            self.is_recording = False
            if self.recording_thread:
                self.recording_thread.join()
            print("音频采集已停止")
    
    async def get_audio_chunk(self):
        """获取音频chunk"""
        return await self.audio_queue.get()
    
    def close(self):
        """清理资源"""
        # 注意：这是一个同步方法，不能调用异步的 stop_recording
        if self.is_recording:
            self.is_recording = False
            if self.recording_thread:
                self.recording_thread.join()
            print("音频采集已停止")
        if self.pyaudio_instance:
            self.pyaudio_instance.terminate()


# 示例使用代码
async def main():
    """示例：启动音频采集并回传"""
    streamer = AudioStreamer()
    
    try:
        # 启动音频采集和回传
        await streamer.start_recording()
        
        print("音频采集和回传已启动，您可以听到麦克风输入的声音...")
        print("按 Ctrl+C 停止")
        
        # 保持运行
        while True:
            await asyncio.sleep(1)
            
    except KeyboardInterrupt:
        print("\n停止音频采集...")
    finally:
        await streamer.stop_recording()
        streamer.close()


if __name__ == "__main__":
    # 运行示例
    asyncio.run(main())